{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic simulation\n",
    "\n",
    "The acrobot is a core example in Drake.  We could certainly load it from a .urdf file, but Drake offers an Acrobot implementation that makes it convenient to manipulate the parameters (and visualize the system with different parameters)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython import get_ipython\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from pydrake.all import (AddMultibodyPlantSceneGraph, ConstantVectorSource, \n",
    "                         DiagramBuilder, Parser, PlanarSceneGraphVisualizer, \n",
    "                         SceneGraph, Simulator)\n",
    "from underactuated import FindResource\n",
    "from underactuated.jupyter import AdvanceToAndVisualize, SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "builder = DiagramBuilder()\n",
    "plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
    "file_name = FindResource(\"models/cartpole.urdf\")\n",
    "Parser(plant).AddModelFromFile(file_name)\n",
    "plant.Finalize()\n",
    "\n",
    "# Set the constant force (you can change the value here if you like)\n",
    "torque_system = builder.AddSystem(ConstantVectorSource([0]))\n",
    "builder.Connect(torque_system.get_output_port(0), plant.get_actuation_input_port())\n",
    "\n",
    "# Setup visualization\n",
    "visualizer = builder.AddSystem(\n",
    "    PlanarSceneGraphVisualizer(scene_graph, xlim=[-2.5, 2.5], ylim=[-1, 2.5], show=plt_is_interactive))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
    "                visualizer.get_input_port(0))\n",
    "\n",
    "diagram = builder.Build()\n",
    "\n",
    "# Set up a simulator to run this diagram\n",
    "simulator = Simulator(diagram)\n",
    "context = simulator.get_mutable_context()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running the cell above, you should see an interactive slider that will control the torque.  Executing the cell below will use the value of that torque as the control input to the pendulum.  Give it a spin!\n",
    "\n",
    "*Note:* If you are running on a system that supports interactive graphics (e.g. a local machine or Binder), then you should see the simulation \"live\", and immediately see the results of changing the slider.  On Colab, however, you only get to see the results of the simulation as a movie after it's run.  [More information.](http://underactuated.mit.edu/drake.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate\n",
    "duration = 10.0 if get_ipython() else 0.1 # sets a shorter duration during testing\n",
    "context.SetTime(0.0)\n",
    "context.SetContinuousState(np.random.randn(4,1))\n",
    "AdvanceToAndVisualize(simulator, visualizer, duration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}